import csv
import logging
import os
import statistics
from collections import defaultdict
from datetime import datetime

import hydra
import torch
from omegaconf import DictConfig, OmegaConf, open_dict

from main import Processor, init_seed

log = logging.getLogger(__name__)


@hydra.main(config_path="config", config_name="experiment/search", version_base=None)
def run_parameter_sweep(cfg: DictConfig):
    """运行参数扫描实验"""

    # 添加时间戳到实验名称
    with open_dict(cfg):
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        cfg.experiment.name = f"{cfg.experiment.name}_{timestamp}"

    all_results = []

    # 打印配置信息
    log.info(f"参数扫描配置:\n{OmegaConf.to_yaml(cfg)}")

    # 遍历参数配置
    for param_idx, param_config in enumerate(cfg.sweep):
        log.info(f"参数配置 {param_idx + 1}/{len(cfg.sweep)}: {param_config}")

        # 更新当前配置
        current_cfg = OmegaConf.create(OmegaConf.to_container(cfg, resolve=True))

        # 在创建处理器前提取model命名空间
        with open_dict(current_cfg):
            # 如果配置被加载到model命名空间，将其提升到顶层
            if "model" in current_cfg:
                for key, value in current_cfg.model.items():
                    current_cfg[key] = value  # 复制到顶层

        for key, value in param_config.items():
            OmegaConf.update(current_cfg, key, value)

        # 配置名称
        config_name = f"config_{param_idx}"
        if "datacase" in param_config:
            config_name = param_config.datacase

        # 更新工作目录
        with open_dict(current_cfg):
            current_cfg.experiment.name = f"{current_cfg.experiment.name}/{config_name}"

        # 运行每个种子
        config_results = []
        for seed_idx, seed in enumerate(current_cfg.experiment.seeds):
            log.info(
                f"  运行种子 {seed_idx + 1}/{len(current_cfg.experiment.seeds)} - 种子: {seed}"
            )

            # 更新种子和工作目录
            run_dir = f"{current_cfg.work_dir}/{config_name}/seed_{seed}"
            with open_dict(current_cfg):
                current_cfg.seed = seed
                current_cfg.work_dir = run_dir

            # 确保目录存在
            os.makedirs(run_dir, exist_ok=True)

            # 重置随机种子
            init_seed(seed)

            # 创建处理器并运行
            processor = Processor(current_cfg)
            processor.start()

            # 记录结果
            result = {
                "config_idx": param_idx,
                "config_name": config_name,
                "seed": seed,
                "accuracy": processor.best_acc,
                "epoch": processor.best_acc_epoch,
            }
            # 添加参数配置
            for key, value in param_config.items():
                result[key] = value

            config_results.append(result)
            all_results.append(result)

            # 清理GPU内存
            torch.cuda.empty_cache()

        # 计算这个配置的统计数据
        accuracies = [r["accuracy"] for r in config_results]
        avg_accuracy = sum(accuracies) / len(accuracies)
        std_accuracy = statistics.stdev(accuracies) if len(accuracies) > 1 else 0.0

        log.info(f"配置 {config_name} 结果: {avg_accuracy:.4f} ± {std_accuracy:.4f}")

    # 创建结果摘要 - 按配置名称分组
    summary = {}
    config_groups = defaultdict(list)

    # 按配置名分组结果
    for result in all_results:
        config_name = result["config_name"]
        config_groups[config_name].append(result)

    # 计算每个配置的汇总统计数据
    for config_name, results in config_groups.items():
        accuracies = [r["accuracy"] for r in results]
        epochs = [r["epoch"] for r in results]

        summary[config_name] = {
            "accuracy_mean": sum(accuracies) / len(accuracies),
            "accuracy_std": statistics.stdev(accuracies)
            if len(accuracies) > 1
            else 0.0,
            "accuracy_count": len(accuracies),
            "epoch_mean": sum(epochs) / len(epochs),
        }

    # 保存结果
    results_dir = f"{cfg.work_dir}/{cfg.experiment.name}"
    os.makedirs(results_dir, exist_ok=True)

    # 保存所有结果到CSV
    with open(f"{results_dir}/all_results.csv", "w", newline="") as f:
        if all_results:
            writer = csv.DictWriter(f, fieldnames=all_results[0].keys())
            writer.writeheader()
            writer.writerows(all_results)

    # 保存摘要到CSV
    with open(f"{results_dir}/summary.csv", "w", newline="") as f:
        fieldnames = [
            "config_name",
            "accuracy_mean",
            "accuracy_std",
            "accuracy_count",
            "epoch_mean",
        ]
        writer = csv.writer(f)
        writer.writerow(fieldnames)
        for config_name, stats in summary.items():
            writer.writerow(
                [
                    config_name,
                    stats["accuracy_mean"],
                    stats["accuracy_std"],
                    stats["accuracy_count"],
                    stats["epoch_mean"],
                ]
            )

    # 输出结果汇总
    log.info("\n参数扫描结果汇总:")
    for config_name, stats in summary.items():
        log.info(
            f"{config_name}: {stats['accuracy_mean']:.4f} ± {stats['accuracy_std']:.4f} ({stats['accuracy_count']} runs)"
        )

    # 找出最佳配置
    if summary:
        best_config_name = max(summary.items(), key=lambda x: x[1]["accuracy_mean"])[0]
        best_stats = summary[best_config_name]
        log.info(f"\n最佳配置: {best_config_name}")
        log.info(
            f"最佳准确率: {best_stats['accuracy_mean']:.4f} ± {best_stats['accuracy_std']:.4f}"
        )

    return summary


if __name__ == "__main__":
    run_parameter_sweep()
